#include "ZNet.h"

#ifdef WIN32
#include <Mstcpip.h>
#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "winmm.lib")
#endif

ZNet::ZNet(int max_service, int the_max_worker, int capacity, int send_size, int recv_size) {
    bService = false;
#ifdef WIN32
    WSADATA wsa_data;
    if(WSAStartup(MAKEWORD(2, 0), &wsa_data ) != 0) return;
#endif
    _send_buf_size = send_size;
    _recv_buf_size = recv_size;
    if(max_service) services = new ZList<ZListenThread>(max_service);
	else services = 0;
    conns = new ZList<ZConn>(capacity);
    max_worker = the_max_worker;
    workers = new ZList<ZNetWorkThread>(max_worker);
    int index;
    for(index = 0; index < max_worker; index++) {
        int i = workers->Add(STATE_WORKING);
        (*workers)[i].max_data_size = send_size;
        (*workers)[i].net = this;
        (*workers)[i].no = index;
#ifdef EPOLL
        (*workers)[i].epoll_size = capacity / max_worker;
#endif
        (*workers)[i].Start();
    }
}

ZNet::~ZNet() {
    StopService();
    delete conns;
    if(services) delete services;
    delete workers;
#ifdef WIN32
    WSACleanup();
#endif
}


bool ZNet::StartService(string local_addr, unsigned short port, bool bCompress) {
    int index = services->Add(STATE_LISTENING);
    if(index >= 0) {
        (*services)[index].net = this;
        (*services)[index].Start(local_addr, port, bCompress);
        bService = true;
        return true;
    }
    return false;
}

void ZNet::StopService() {
	int index;
    if(bService) {
		for(index = 0; index < services->max_number; index++) if((*services)[index].state != STATE_NONE) (*services)[index].Stop();
	}
	for(index = 0; index < workers->max_number; index++) if((*workers)[index].state != STATE_NONE) (*workers)[index].Stop();
	for(index = 0; index < conns->max_number; index++) if((*conns)[index].state == STATE_IDLE) (*conns)[index].Shutdown();
	ZThread::Sleep(100);
}

void ZNet::PrepareConn(int index) {
        ZConn &conn = (*conns)[index];
#ifdef WIN32
        char opt = 1;
        setsockopt((unsigned int)conn.sock, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(char));
#else
        fcntl(s, F_SETFL, O_NONBLOCK);
#endif
        conn.Init(conn.sock, _send_buf_size, _recv_buf_size, conn.bCompress);
#ifdef WIN32
        conn.hIOCP = CreateIoCompletionPort((HANDLE)conn.sock, (*workers)[index % workers->capacity].hIOCP, (ULONG_PTR)&conn, 1);
#endif
#ifdef EPOLL
        struct epoll_event event;
        event.data.ptr = (void *)&conn;
        event.events = EPOLLIN | EPOLLET;
        epoll_ctl((*workers)[index % workers->capacity].epoll_fd, EPOLL_CTL_ADD, conn.sock, &event);
#endif
        conn.Recv();
}

int ZNet::AddConn(SOCKET s, sockaddr_in* addr, bool bCompress) {
    int index = conns->Add(STATE_CONNECTING);
    if(index >= 0) {
        ZConn &conn = (*conns)[index];
		conn.bCompress = bCompress;
        conn.sock = s;
        memcpy(&conn.remote, addr, sizeof(sockaddr_in));
        int length = sizeof(sockaddr_in);
        getsockname(conn.sock, (sockaddr*)&conn.local, &length);
        PrepareConn(index);
        conn.state = STATE_CONNECTED;
    }
    return index;
}

int ZNet::Connect(string remote_addr, unsigned short port, bool bCompress) {
    struct sockaddr_in address;
    memset((void*)&address, 0, sizeof(address));
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = inet_addr(remote_addr.c_str());
    address.sin_port = htons(port);
    SOCKET s = socket(AF_INET, SOCK_STREAM, 0);
    if(connect(s, (sockaddr *)&address, sizeof(address)) < 0) {
        closesocket(s);
        return -1;
    }
    else return AddConn(s, &address, bCompress);
}

void ZListenThread::Run() {
    if(listen(listen_sock, 64) >= 0) {
        while(!bStop) {
            SOCKET s;
            struct sockaddr_in client_addr;
            socklen_t len = sizeof(client_addr);
            memset(&client_addr, 0, sizeof(client_addr));
            s = accept(listen_sock, (sockaddr*)&client_addr, &len);
            if(!bStop && s > 0) net->AddConn(s, &client_addr, bCompress);
        }
    }
}

void ZConn::Send() {
    ZBuffer* buf = send_buf;
    if(!buf) buf = SendBuf;
    if(!bSending && buf->Length() > 0) {
        bSending = true;
#ifdef WIN32
        UINT nRetVal;
        DWORD dwIoSize;
        ULONG ulFlags = MSG_PARTIAL;
        WSABUF wsabuf[1];
		unsigned int length;
        wsabuf[0].buf = buf->GetData(length);
		wsabuf[0].len = (ULONG)length;
        send_overlap._pos = buf->_start_pos;
        send_overlap._size = wsabuf[0].len;
        nRetVal = WSASend(sock, wsabuf, 1, &dwIoSize, ulFlags, &send_overlap.ol, NULL);
        if(nRetVal == SOCKET_ERROR) {
            if(WSAGetLastError() != WSA_IO_PENDING) Shutdown();
        }
#endif
#ifdef EPOLL
        unsigned long size;
        char *buf = buf->GetData(size);
        int length = send(sock, buf, size, 0);
        if(length <= 0) {
            if(SOCKET_ERR() != EWOULDBLOCK) Shutdown();
        }
        else {
            buf->Remove(length);
            bSending = false;
            Send();
        }
#endif
    }
}

void ZConn::Recv() {
    ZBuffer* buf = recv_buf;
    if(!buf) buf = RecvBuf;
    if(buf->Length() < buf->capacity - (int)sizeof(int)) {
#ifdef WIN32
        UINT nRetVal;
        DWORD dwIoSize;
        ULONG ulFlags;
        WSABUF wsabuf[1];
		unsigned int length;
        wsabuf[0].buf = buf->GetBuf(length);
		wsabuf[0].len = (ULONG)length;
        ulFlags = MSG_PARTIAL;
        dwIoSize = 0;
        recv_overlap._pos = buf->_stop_pos;
        recv_overlap._size = wsabuf[0].len;
        nRetVal = WSARecv(sock, wsabuf, 1, &dwIoSize, &ulFlags, &recv_overlap.ol, NULL);
        if(nRetVal == SOCKET_ERROR) {
            if(WSAGetLastError() != WSA_IO_PENDING)Shutdown();
        }
#endif
#ifdef EPOLL
        unsigned long size;
        char *buf = buf->GetBuf(size);
        long length = recv(sock, buf, size, 0);
        if(length < 0 && SOCKET_ERR() != EWOULDBLOCK) Shutdown();
        else if(length > 0) {
            buf->Add(length);
            Recv();
        }
#endif
    }
    else state = STATE_CLOSING;
}
